# -*- coding: utf-8 -*-
# @Author  : longbhu
# @Time    : 2025/3/5 14:59
# @Function: 同过土壤保持的实物量来计算价值量
"""
    通过土壤保持的实物量来计算价值量。
输入：
    Qsr：土壤保持的实物量：某文件路径       数据类型：tif
    泥沙淤积系数：0.24        数据类型：常量
    土壤容重：1.27        数据类型：常量
    水库单位清淤工程费用：60        数据类型：常量
    sr_TN: 污染物氮纯含量：某文件路径     数据类型：tif
    sr_TP: 污染物磷纯含量：某文件路径     数据类型：tif
    污染物氮单位处理成本：3000        数据类型：常量
    污染物磷单位处理成本：3500        数据类型：常量
输出：
    Vsr：土壤保持价值      数据类型：tif        计算方法：Vsr = Vsd + Vdpd
    Vsd: 减少泥沙淤积价值       数据类型：tif        计算方法：Vsd = 泥沙淤积系数 * Qsr / 土壤容重 * 水库单位清淤工程费用
    Vdpd: 减少面源污染价值      数据类型：tif        计算方法：Vdpd = Qsr *（ sr_TN * 污染物氮单位处理成本 + sr_TP * 污染物磷单位处理成本）
"""
# 数据的前期准备
# 根据shp裁剪保证Qsr与sr_TN和sr_TP的大小一致

from rasterio.warp import reproject
import rasterio
from rasterio.enums import Resampling
import numpy as np

# todo 应设置统一的输出路径，在核算时检查是否有应该输出的tif文件，如果有则直接统计结果；否则需要输出tif文件后再核算统计结果。

def resample_to_match(target_src, source_src):
    """将源栅格重采样至目标栅格的坐标系和尺寸"""
    destination = np.zeros(target_src.shape, dtype=np.float32)

    # 处理源数据无效值
    source_data = source_src.read(1)
    source_nodata = source_src.nodata
    source_data = np.where(source_data == source_nodata, np.nan, source_data)

    reproject(
        source=source_data,
        destination=destination,
        src_transform=source_src.transform,
        src_crs=source_src.crs,
        dst_transform=target_src.transform,
        dst_crs=target_src.crs,
        resampling=Resampling.nearest
    )
    return destination


def calculate_soil_conservation_value(Qsr_path, sr_TN_path, sr_TP_path, mud_siltation_coefficient=0.24, soil_bulk_density=1.27, reservoir_siltation_cost=60, TN_treatment_cost=3000, TP_treatment_cost=3500):
    """
    计算土壤保持价值。

    参数:
    Qsr_path (str): 土壤保持的实物量文件路径。
    sr_TN_path (str): 污染物氮纯含量文件路径。
    sr_TP_path (str): 污染物磷纯含量文件路径。
    mud_siltation_coefficient (float): 泥沙淤积系数。默认为0.24。
    soil_bulk_density (float): 土壤容重。默认为1.27。
    reservoir_siltation_cost (float): 水库单位清淤工程费用。默认为60。
    TN_treatment_cost (float): 污染物氮单位处理成本。默认为3000。
    TP_treatment_cost (float): 污染物磷单位处理成本。默认为3500。

    返回:
    Vsd (numpy.ndarray): 减少泥沙淤积价值。
    Vdpd (numpy.ndarray): 减少面源污染价值。
    Vsr (numpy.ndarray): 土壤保持价值。
    """
    with rasterio.open(Qsr_path) as src_Qsr:
        Qsr = src_Qsr.read(1)
        qsr_nodata = src_Qsr.nodata
        Qsr = np.where(Qsr == qsr_nodata, np.nan, Qsr)
        qsr_meta = src_Qsr.meta

    with rasterio.open(sr_TN_path) as src_sr_TN:
        sr_TN = resample_to_match(src_Qsr, src_sr_TN)

    with rasterio.open(sr_TP_path) as src_sr_TP:
        sr_TP = resample_to_match(src_Qsr, src_sr_TP)

    # 计算时自动排除nan值
    Vsd = mud_siltation_coefficient * Qsr / soil_bulk_density * reservoir_siltation_cost
    Vdpd = Qsr * (sr_TN * TN_treatment_cost / 1000 / 100 + sr_TP * TP_treatment_cost / 1000 / 100)  # / 1000是sr_TN的数据转换 / 100是百分比
    Vsr = Vsd + Vdpd

    return Vsd, Vdpd, Vsr


def save_array_as_tiff(array, output_path, src, nodata=-9999):
    array = np.where(np.isnan(array), nodata, array)

    meta = src.meta.copy()
    meta.update({
        'driver': 'GTiff',
        'count': 1,
        'dtype': array.dtype,
        'nodata': nodata
    })

    with rasterio.open(output_path, 'w', **meta) as dst:
        dst.write(array, 1)

# 示例用法
Qsr_path = r'E:\GEP_data\sr\Qsr\lingchuan_merged_Qsr.tif'
sr_TN_path =r'E:\GEP_data\sr\lingchuan_merged_tn_Pro.tif'
sr_TP_path = r'E:\GEP_data\sr\lingchuan_merged_tp_Pro.tif'

with rasterio.open(Qsr_path) as src:
    Vsd, Vdpd, Vsr = calculate_soil_conservation_value(Qsr_path, sr_TN_path, sr_TP_path)

    # 保存时指定nodata值（与原始Qsr保持一致）
    save_array_as_tiff(Vsd, 'Vsd.tif', src, nodata=src.nodata)
    save_array_as_tiff(Vdpd, 'Vdpd.tif', src, nodata=src.nodata)
    save_array_as_tiff(Vsr, 'Vsr.tif', src, nodata=src.nodata)

